package com.aisino.encrypt.algorithm;

import com.aisino.encrypt.exception.AisinoException;

import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;
import javax.crypto.*;
import javax.crypto.interfaces.DHPrivateKey;
import javax.crypto.interfaces.DHPublicKey;
import javax.crypto.spec.DHParameterSpec;

/**
 * DH加密
 *
 * @author 蒋帆
 * @datetime 2021-08-27
 * @description 1.甲方构建密钥对儿，将公钥公布给乙方，将私钥保留；双方约定数据加密算法；乙方通过甲方公钥构建密钥对儿，将公钥公布给甲方，将私钥保留。
 * 2.甲方使用私钥、乙方公钥、约定数据加密算法构建本地密钥，然后通过本地密钥加密数据，发送给乙方加密后的数据；乙方使用私钥、甲方公钥、约定数据加密算法构建本地密钥，然后通过本地密钥对数据解密。
 * 3.乙方使用私钥、甲方公钥、约定数据加密算法构建本地密钥，然后通过本地密钥加密数据，发送给甲方加密后的数据；甲方使用私钥、乙方公钥、约定数据加密算法构建本地密钥，然后通过本地密钥对数据解密。
 */
public class DH {
    /**
     * 公钥
     */
    public static final String PUBLIC_KEY = "DHPublicKey";
    /**
     * 私钥
     */
    public static final String PRIVATE_KEY = "DHPrivateKey";
    /**
     * 开始生成本地密钥SecretKey 密钥算法为对称密码算法
     * 可以为 DES DES AES
     */
    public static final String KEY_ALGORITHM_DES = "DES";

    /**
     * 定义加密方式
     */
    private static final String KEY_ALGORITHM = "DH";

    /**
     * 密钥长度，DH算法的默认密钥长度是1024
     * 密钥长度必须是64的倍数，在512到65536位之间
     */
    private static final int KEY_SIZE = 1024;

    /**
     * 编码
     */
    private static final int INTEGER_HEX = 0xFF;

    /**
     * 1
     */
    private static final int ONE = 1;

    /**
     * 甲方初始化并返回密钥对
     *
     * @return 返回生成的密钥对
     */
    public static Map<String, Object> initKey() {
        Map<String, Object> keyMap = null;
        try {
            // 实例化密钥对生成器
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(KEY_ALGORITHM);
            // 初始化密钥对生成器 默认是1024 512-1024 & 64的倍数
            keyPairGenerator.initialize(KEY_SIZE);
            // 生成密钥对
            KeyPair keyPair = keyPairGenerator.generateKeyPair();
            // 得到甲方公钥
            DHPublicKey publicKey = (DHPublicKey) keyPair.getPublic();
            // 得到甲方私钥
            DHPrivateKey privateKey = (DHPrivateKey) keyPair.getPrivate();
            // 将公钥和私钥封装在Map中， 方便之后使用
            keyMap = new HashMap<String, Object>();
            keyMap.put(PUBLIC_KEY, publicKey);
            keyMap.put(PRIVATE_KEY, privateKey);
        } catch (NoSuchAlgorithmException e) {
            throw new AisinoException(e);
        }
        return keyMap;
    }

    /**
     * 传入密钥大小甲方初始化并返回密钥对
     *
     * @param keySize 初始化密钥对生成器，密钥大小为64的倍数，在512到65536位之间
     * @return 返回生成的密钥对
     */
    public static Map<String, Object> initKey(int keySize) {
        if (keySize < 512 || keySize > 65536) {
            throw new AisinoException("密钥大小应为64的倍数，且在512到65536位之间");
        }
        Map<String, Object> keyMap = null;
        try {
            // 实例化密钥对生成器
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(KEY_ALGORITHM);
            // 初始化密钥对生成器 默认是1024 512-1024 & 64的倍数
            keyPairGenerator.initialize(keySize);
            // 生成密钥对
            KeyPair keyPair = keyPairGenerator.generateKeyPair();
            // 得到甲方公钥
            DHPublicKey publicKey = (DHPublicKey) keyPair.getPublic();
            // 得到甲方私钥
            DHPrivateKey privateKey = (DHPrivateKey) keyPair.getPrivate();
            // 将公钥和私钥封装在Map中， 方便之后使用
            keyMap = new HashMap<String, Object>();
            keyMap.put(PUBLIC_KEY, publicKey);
            keyMap.put(PRIVATE_KEY, privateKey);
        } catch (NoSuchAlgorithmException e) {
            throw new AisinoException(e);
        }
        return keyMap;
    }

    /**
     * 乙方根据甲方公钥初始化并返回密钥对
     *
     * @param key 甲方的公钥
     * @return 返回乙方生成的密钥对
     */
    public static Map<String, Object> initKey(String key) {
        try {
            // 将甲方公钥从字节数组转换为PublicKey
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(Base64.getDecoder().decode(key));
            // 实例化密钥工厂
            KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
            // 产生甲方公钥pubKey
            DHPublicKey dhPublicKey = (DHPublicKey) keyFactory
                    .generatePublic(keySpec);
            // 剖析甲方公钥，得到其参数
            DHParameterSpec dhParameterSpec = dhPublicKey.getParams();
            // 实例化密钥对生成器
            KeyPairGenerator keyPairGenerator = KeyPairGenerator
                    .getInstance(KEY_ALGORITHM);
            // 用甲方公钥初始化密钥对生成器
            keyPairGenerator.initialize(dhParameterSpec);
            // 产生密钥对
            KeyPair keyPair = keyPairGenerator.generateKeyPair();
            // 得到乙方公钥
            DHPublicKey publicKey = (DHPublicKey) keyPair.getPublic();
            // 得到乙方私钥
            DHPrivateKey privateKey = (DHPrivateKey) keyPair.getPrivate();
            // 将公钥和私钥封装在Map中， 方便之后使用
            Map<String, Object> keyMap = new HashMap<String, Object>();
            keyMap.put(PUBLIC_KEY, publicKey);
            keyMap.put(PRIVATE_KEY, privateKey);
            return keyMap;
        } catch (InvalidAlgorithmParameterException e) {
            throw new AisinoException(e);
        } catch (NoSuchAlgorithmException e) {
            throw new AisinoException(e);
        } catch (InvalidKeySpecException e) {
            throw new AisinoException(e);
        }
    }

    /**
     * 根据对方的公钥和自己的私钥生成 本地密钥,返回的是SecretKey对象的字节数组
     *
     * @param publicKey  公钥
     * @param privateKey 私钥
     * @return 返回SecretKey对象的字节数组
     */
    public static String getSecretKeyBytes(String publicKey, String privateKey) {
        try {
            // 实例化密钥工厂
            KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
            // 将公钥从字节数组转换为PublicKey
            X509EncodedKeySpec pubKeySpec = new X509EncodedKeySpec(Base64.getDecoder().decode(publicKey));
            PublicKey pubKey = keyFactory.generatePublic(pubKeySpec);
            // 将私钥从字节数组转换为PrivateKey
            PKCS8EncodedKeySpec priKeySpec = new PKCS8EncodedKeySpec(Base64.getDecoder().decode(privateKey));
            PrivateKey priKey = keyFactory.generatePrivate(priKeySpec);

            // 准备根据以上公钥和私钥生成本地密钥SecretKey
            // 先实例化KeyAgreement
            KeyAgreement keyAgreement = KeyAgreement.getInstance(KEY_ALGORITHM);
            // 用自己的私钥初始化keyAgreement
            keyAgreement.init(priKey);
            // 结合对方的公钥进行运算
            keyAgreement.doPhase(pubKey, true);
            // 开始生成本地密钥SecretKey 密钥算法为对称密码算法
            SecretKey secretKey = keyAgreement.generateSecret(KEY_ALGORITHM_DES);
            return Base64.getEncoder().encodeToString(secretKey.getEncoded());
        } catch (NoSuchAlgorithmException e) {
            throw new AisinoException(e);
        } catch (InvalidKeySpecException e) {
            throw new AisinoException(e);
        } catch (InvalidKeyException e) {
            throw new AisinoException(e);
        }

    }

    /**
     * 根据对方的公钥和自己的私钥生成 本地密钥,返回的是SecretKey对象
     *
     * @param publicKey  公钥
     * @param privateKey 私钥
     * @return 返回的是SecretKey对象
     */
    private static SecretKey getSecretKey(String publicKey, String privateKey) {
        SecretKey secretKey = null;
        try {
            // 实例化密钥工厂
            KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
            // 将公钥从字节数组转换为PublicKey
            X509EncodedKeySpec pubKeySpec = new X509EncodedKeySpec(Base64.getDecoder().decode(publicKey));
            PublicKey pubKey = keyFactory.generatePublic(pubKeySpec);
            // 将私钥从字节数组转换为PrivateKey
            PKCS8EncodedKeySpec priKeySpec = new PKCS8EncodedKeySpec(Base64.getDecoder().decode(privateKey));
            PrivateKey priKey = keyFactory.generatePrivate(priKeySpec);

            // 准备根据以上公钥和私钥生成本地密钥SecretKey
            // 先实例化KeyAgreement
            KeyAgreement keyAgreement = KeyAgreement.getInstance(KEY_ALGORITHM);
            // 用自己的私钥初始化keyAgreement
            keyAgreement.init(priKey);
            // 结合对方的公钥进行运算
            keyAgreement.doPhase(pubKey, true);
            // 开始生成本地密钥SecretKey 密钥算法为对称密码算法
            secretKey = keyAgreement.generateSecret(KEY_ALGORITHM_DES);
        } catch (NoSuchAlgorithmException e) {
            throw new AisinoException(e);
        } catch (InvalidKeySpecException e) {
            throw new AisinoException(e);
        } catch (InvalidKeyException e) {
            throw new AisinoException(e);
        }
        return secretKey;
    }

    /**
     * 从 Map 中取得公钥
     *
     * @param keyMap 密钥对
     * @return 返回从 Map 中取得公钥
     */
    public static String getPublicKey(Map<String, Object> keyMap) {
        DHPublicKey key = (DHPublicKey) keyMap.get(PUBLIC_KEY);
        return Base64.getEncoder().encodeToString(key.getEncoded());
    }

    /**
     * 从 Map 中取得私钥
     *
     * @param keyMap 密钥对
     * @return 返回从 Map 中取得的私钥
     */
    public static String getPrivateKey(Map<String, Object> keyMap) {
        DHPrivateKey key = (DHPrivateKey) keyMap.get(PRIVATE_KEY);
        return Base64.getEncoder().encodeToString(key.getEncoded());
    }

    /**
     * DH 加密
     *
     * @param data       带加密数据
     * @param publicKey  甲方公钥
     * @param privateKey 乙方私钥
     * @return 密文
     */
    public static String encryptDH(String data, String publicKey,
                                   String privateKey) {
        byte[] bytes = null;
        try {
            //
            SecretKey secretKey = getSecretKey(publicKey, privateKey);
            // 数据加密
            Cipher cipher = Cipher.getInstance(secretKey.getAlgorithm());
            cipher.init(Cipher.ENCRYPT_MODE, secretKey);
            bytes = cipher.doFinal(data.getBytes());
        } catch (NoSuchPaddingException e) {
            throw new AisinoException(e);
        } catch (IllegalBlockSizeException e) {
            throw new AisinoException(e);
        } catch (NoSuchAlgorithmException e) {
            throw new AisinoException(e);
        } catch (BadPaddingException e) {
            throw new AisinoException(e);
        } catch (InvalidKeyException e) {
            throw new AisinoException(e);
        }
        return Base64.getEncoder().encodeToString(bytes);
    }

    /**
     * DH 解密
     *
     * @param data       待解密数据
     * @param publicKey  乙方公钥
     * @param privateKey 甲方私钥
     * @return 解密后的内容
     */
    public static String decryptDH(String data, String publicKey,
                                   String privateKey) {
        byte[] bytes = null;
        try {
            //
            SecretKey secretKey = getSecretKey(publicKey, privateKey);
            // 数据加密
            Cipher cipher = Cipher.getInstance(secretKey.getAlgorithm());
            cipher.init(Cipher.DECRYPT_MODE, secretKey);
            bytes = cipher.doFinal(Base64.getDecoder().decode(data));
        } catch (NoSuchPaddingException e) {
            throw new AisinoException(e);
        } catch (IllegalBlockSizeException e) {
            throw new AisinoException(e);
        } catch (NoSuchAlgorithmException e) {
            throw new AisinoException(e);
        } catch (BadPaddingException e) {
            throw new AisinoException(e);
        } catch (InvalidKeyException e) {
            throw new AisinoException(e);
        }
        return new String(bytes);
    }

    /**
     * 字节数组转16进制
     *
     * @param resultBytes 字节数组
     * @return 字节数组转16进制
     */
    public static String fromBytesToHex(byte[] resultBytes) {
        StringBuilder builder = new StringBuilder();
        for (byte b : resultBytes
        ) {
            if (Integer.toHexString(INTEGER_HEX & b).length() == ONE) {
                builder.append("0").append(
                        Integer.toHexString(INTEGER_HEX & b));
            } else {
                builder.append(Integer.toHexString(INTEGER_HEX & b));
            }
        }
        return builder.toString();
    }
}
